# main.py
"""
项目主入口文件。
负责初始化服务、协调数据流并执行整个数据集生成流程。
"""
import asyncio
import logging
import os
import nest_asyncio

from llama_index.core import Settings
from llama_index.core.embeddings import resolve_embed_model
from llama_index.llms.deepseek import DeepSeek

# 导入配置和模块
import config
import data_fetcher
import dataset_generator
import utils

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def setup_services():
    """初始化并配置所有外部服务，如 LLM 和嵌入模型。"""
    logging.info("正在初始化服务...")

    # 设置 API 密钥
    os.environ["TAVILY_API_KEY"] = config.TAVILY_API_KEY
    os.environ['DEEPSEEK_API_KEY'] = config.DEEPSEEK_API_KEY

    # 初始化 LLM 和嵌入模型
    llm = DeepSeek(model=config.LLM_MODEL_NAME)
    embed_model = resolve_embed_model(config.EMBED_MODEL_PATH)

    # 应用到 LlamaIndex 的全局设置
    Settings.llm = llm
    Settings.embed_model = embed_model

    logging.info("服务初始化完成。")
    return llm


async def main():
    """主执行函数"""
    # nest_asyncio.apply() 允许在已经有事件循环的环境（如Jupyter）中运行asyncio
    # 在纯 .py 脚本中，如果确定没有其他asyncio循环，可以不加，但加上能提高兼容性
    nest_asyncio.apply()

    # 1. 初始化
    llm = setup_services()

    try:
        # 2. 获取数据
        documents = data_fetcher.fetch_documents_from_tavily(
            api_key=config.TAVILY_API_KEY,
            query=config.SEARCH_QUERY,
            max_results=config.MAX_SEARCH_RESULTS
        )
        if not documents:
            logging.warning("未能获取到任何文档，程序提前结束。")
            return

        # 3. 生成 RAG 数据集
        rag_dataset = await dataset_generator.generate_rag_dataset(
            documents=documents,
            llm=llm,
            chunk_size=config.CHUNK_SIZE,
            chunk_overlap=config.CHUNK_OVERLAP,
            questions_per_chunk=config.NUM_QUESTIONS_PER_CHUNK
        )

        # 保存原始的 RagDataset 对象，方便后续分析
        rag_dataset.save_json(config.RAW_RAG_DATASET_PATH)
        logging.info(f"原始 RagDataset 已保存到 {config.RAW_RAG_DATASET_PATH}")

        # 4. 格式化并保存 SFT 数据集
        sft_data = utils.format_rag_for_sfttrainer(rag_dataset)
        utils.save_json(sft_data, config.SFT_DATASET_PATH)

        logging.info("所有流程执行完毕！")

    except Exception as e:
        logging.critical(f"程序执行过程中发生未处理的异常: {e}", exc_info=True)


if __name__ == "__main__":
    # 使用 asyncio.run() 来执行异步的 main 函数
    asyncio.run(main())